#include <gtest/gtest.h>
#include "Thread/semaphore.h"
#include "Thread/threadgroup.h"
#include "Util/TimeTicker.h"
#include <queue>
#include <functional>
#include <atomic>
#include <mutex>

using namespace toolkit;

class TaskQueue {
public:
    TaskQueue(size_t num_workers = 4) : _stop(false) {
        for (size_t i = 0; i < num_workers; ++i) {
            _workers.create_thread([this]() { worker(); });
        }
    }
    
    ~TaskQueue() {
        _stop = true;
        _task_sem.post(_workers.size());
        _workers.join_all();
    }
    
    void addTask(std::function<void()> task) {
        std::lock_guard<std::mutex> lock(_mutex);
        _tasks.push(task);
        _task_sem.post();
    }
    
private:
    void worker() {
        while (!_stop) {
            _task_sem.wait();
            if (_stop) break;
            
            std::function<void()> task;
            {
                std::lock_guard<std::mutex> lock(_mutex);
                if (!_tasks.empty()) {
                    task = _tasks.front();
                    _tasks.pop();
                }
            }
            
            if (task) {
                task();
            }
        }
    }
    
    semaphore _task_sem;
    threadgroup _workers;
    std::queue<std::function<void()>> _tasks;
    std::mutex _mutex;
    std::atomic<bool> _stop;
};

class TaskQueueTest : public ::testing::Test {
protected:
    void SetUp() override {}
    void TearDown() override {}
};

// 基本功能测试
TEST_F(TaskQueueTest, BasicFunctionality) {
    TaskQueue queue(2);
    std::atomic<int> counter(0);
    
    for (int i = 0; i < 10; ++i) {
        queue.addTask([&]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            counter++;
        });
    }
    
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    EXPECT_EQ(counter, 10);
}

// 并发测试
TEST_F(TaskQueueTest, ConcurrentAccess) {
    TaskQueue queue(4);
    std::atomic<int> counter(0);
    const int NUM_TASKS = 1000;
    
    TimeTicker ticker;
    
    for (int i = 0; i < NUM_TASKS; ++i) {
        queue.addTask([&]() {
            counter++;
        });
    }
    
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    EXPECT_EQ(counter, NUM_TASKS);
    
    std::cout << "Processed " << NUM_TASKS << " tasks in " 
              << ticker.elapsedTime() << "ms" << std::endl;
}

// 任务优先级测试
TEST_F(TaskQueueTest, TaskPriority) {
    TaskQueue queue(1);  // 单线程执行以确保顺序
    std::vector<int> execution_order;
    std::mutex order_mutex;
    
    // 添加低优先级任务
    for (int i = 0; i < 5; ++i) {
        queue.addTask([&, i]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            std::lock_guard<std::mutex> lock(order_mutex);
            execution_order.push_back(i);
        });
    }
    
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    
    // 验证执行顺序
    for (size_t i = 0; i < execution_order.size(); ++i) {
        EXPECT_EQ(execution_order[i], i);
    }
}

// 压力测试
TEST_F(TaskQueueTest, StressTest) {
    TaskQueue queue(4);
    std::atomic<int> counter(0);
    const int NUM_TASKS = 10000;
    
    TimeTicker ticker;
    
    // 创建多个生产者线程
    std::vector<std::thread> producers;
    for (int i = 0; i < 4; ++i) {
        producers.emplace_back([&]() {
            for (int j = 0; j < NUM_TASKS/4; ++j) {
                queue.addTask([&]() {
                    counter++;
                });
            }
        });
    }
    
    for (auto& t : producers) {
        t.join();
    }
    
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    EXPECT_EQ(counter, NUM_TASKS);
    
    std::cout << "Stress test completed in " << ticker.elapsedTime() 
              << "ms" << std::endl;
}

// 异常处理测试
TEST_F(TaskQueueTest, ExceptionHandling) {
    TaskQueue queue(2);
    std::atomic<int> success_count(0);
    std::atomic<int> failure_count(0);
    
    for (int i = 0; i < 10; ++i) {
        queue.addTask([&, i]() {
            try {
                if (i % 2 == 0) {
                    throw std::runtime_error("Test exception");
                }
                success_count++;
            } catch (...) {
                failure_count++;
            }
        });
    }
    
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    EXPECT_EQ(success_count, 5);
    EXPECT_EQ(failure_count, 5);
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
} 